package middleware

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"github.com/go-redis/redis_rate/v9"
	"github.com/vueadmin/utils/response"
	"go.uber.org/zap"
	"net/http"
)

type LimiterMiddleware struct {
	redis *redis.Client
	log   *zap.SugaredLogger
}

func NewLimiterMiddleware(redis *redis.Client, log *zap.SugaredLogger) *LimiterMiddleware {
	return &LimiterMiddleware{
		redis: redis,
		log:   log,
	}
}

// 限流器
func (p LimiterMiddleware) Handler(appname string, limit int) gin.HandlerFunc {
	return func(c *gin.Context) {
		limiter := redis_rate.NewLimiter(p.redis)
		res, err := limiter.Allow(context.Background(), appname, redis_rate.PerSecond(limit))
		if err != nil {
			p.log.Error("限流器错误:", err)
			return
		}
		if res.Remaining == 0 {
			response.FailWithCode(c, http.StatusServiceUnavailable, "请求过于频繁，请稍后重试")
			c.Abort()
			return
		}
		c.Next()
	}
}
